import matplotlib.pyplot as plt
import os
import glob
from PIL import Image


file_path = "D:\PytorchLearn\cornell\data-1"

grasp_file = glob.glob(os.path.join(file_path,"pcd*cpos.txt"))
image_file = glob.glob(os.path.join(file_path,"pcd*r.png"))

image_file.sort()
grasp_file.sort()

count = 0
while (count < len(grasp_file)):
    with open(grasp_file[count]) as f:
        while(True):
            p0 = f.readline()
            if not p0: break
            p1 =f.readline()
            p2 =f.readline()
            p3 =f.readline()

            p0=p0.split()
            p1=p1.split()
            p2=p2.split()
            p3=p3.split()

            p0 = [float(num) for num in p0]
            p1 = [float(num) for num in p1]
            p2 = [float(num) for num in p2]
            p3 = [float(num) for num in p3]

            img = Image.open(image_file[count])
            width, height = img.size
            plt.imshow(img)

            ax = plt.gca()
            ax.set_ylim(0, height)
            ax.set_xlim(0, width)

            plt.plot([p0[0], p1[0]],[p0[1],p1[1]],color='red')
            plt.plot([p1[0], p2[0]],[p1[1],p2[1]],color='red')
            plt.plot([p2[0], p3[0]],[p2[1],p3[1]],color='red')
            plt.plot([p3[0], p0[0]],[p3[1],p0[1]], p0,color='red')
        plt.savefig(f"D:\\PytorchLearn\\cornell\\rec\\{count}.png")
        plt.close()
    count += 1